import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from matplotlib.colors import LinearSegmentedColormap


def plot_heatmap(frequency, begin_date, x_name, y_name):
    for result_label in ["annualized_return", "annualized_volatility", "annualized_shape", "turnover", "margin",
                         "fitness", "excess_return", "relative_return", "max_drawdown", "information_ratio",
                         "calmar_ratio", "yearly_count", "winning_rate", "profit_loss_ratio", "kelly_fraction"]:
        df = pd.read_csv(f'{frequency}/{begin_date}/result/backtest_result_heatmap.csv')
        # 使用 pivot_table 转换数据
        pivot_df = df.pivot_table(
            values=f'{result_label}',
            index=f'{y_name}',      # Y轴作为行
            columns=f'{x_name}',    # X轴作为列
            aggfunc='mean'      # 处理重复值：取平均值
        )

        print("转换后的数据:")
        print(pivot_df)

        # 绘制热力图
        plt.figure(figsize=(10, 8))
        colors = ["green", "yellow", "red"]
        if result_label == "max_drawdown":
            colors = ["green", "yellow", "red"][::-1]
        cmap = LinearSegmentedColormap.from_list("rg", colors, N=256)
        sns.heatmap(
            pivot_df,
            annot=True,           # 显示数值
            fmt='.1f',           # 数值格式
            cmap=cmap,      # 颜色映射
            cbar_kws={'label': 'Value Scale'},
            linewidths=0.5,
            linecolor='white'
        )

        plt.title(f'{result_label}')
        plt.xlabel(f'{x_name}')
        plt.ylabel(f'{y_name}')
        plt.tight_layout()
        # plt.show()
        plt.savefig(f'{frequency}/{begin_date}/heatmap_png/{result_label}.png')